#include <stdio.h>
#include <math.h>
#include <stdbool.h>
#include <string.h>

#define NUM_WAVEPOINTS 400
#define WAVELENGTH_MIN 380
#define MAX_LEDS 10
#define MAX_ITER 1000
#define LEARNING_RATE 0.01f
#define R2_THRESHOLD 0.85f
#define WEIGHT_THRESHOLD 0.01f
#define GRAD_CLIP 1.0f
#define EPSILON 1e-10f

typedef struct {
    float peak_wavelength;
    float fwhm;
    float intensity;
    float sigma;
    float cost; // 新增成本参数
} LEDModel;

typedef struct {
    float target[NUM_WAVEPOINTS];
    LEDModel leds[MAX_LEDS];
    float led_matrix[MAX_LEDS][NUM_WAVEPOINTS];
    float optimal_weights[MAX_LEDS];
    int num_leds;
} SpectrumFitter;

inline void precompute_sigma(LEDModel* led) {
    led->sigma = led->fwhm / (2.0f * sqrtf(2.0f * (float)log(2.0)));
}

inline float gaussian_spectrum(float wavelength, const LEDModel* led) {
    float delta = wavelength - led->peak_wavelength;
    float exponent = -(delta * delta) / (2.0f * led->sigma * led->sigma);
    return led->intensity * expf(exponent);
}

void load_AM15G(float target[]) {
    const float bands[][3] = {{400,50,0.8}, {550,100,1.2}, {700,80,1.0}};
    const int num_bands = sizeof(bands)/sizeof(bands[0]);
    
    memset(target, 0, NUM_WAVEPOINTS*sizeof(float));
    for (int b = 0; b < num_bands; b++) {
        float center = bands[b][0], width = bands[b][1], amp = bands[b][2];
        for (int i = 0; i < NUM_WAVEPOINTS; i++) {
            float wl = WAVELENGTH_MIN + i;
            target[i] += amp * expf(-powf((wl-center)/width, 2));
        }
    }
}

void init_LED_database(LEDModel leds[], int* num_leds) {
    const float params[][4] = {
        {385,15,0.8,50}, {450,25,1.5,60}, {520,30,1.4,70},
        {590,20,1.0,55}, {660,20,1.1,65}, {730,25,0.6,45},
        {450,120,2.0,40}
    };
    
    *num_leds = sizeof(params)/sizeof(params[0]);
    for (int i = 0; i < *num_leds; i++) {
        leds[i] = (LEDModel){
            .peak_wavelength = params[i][0],
            .fwhm = params[i][1],
            .intensity = params[i][2],
            .cost = params[i][3]
        };
        precompute_sigma(&leds[i]);
    }
}

void init_spectrum_fitter(SpectrumFitter* fitter) {
    load_AM15G(fitter->target);
    init_LED_database(fitter->leds, &fitter->num_leds);
    
    for (int j = 0; j < NUM_WAVEPOINTS; j++) {
        float wl = WAVELENGTH_MIN + j;
        for (int i = 0; i < fitter->num_leds; i++) {
            fitter->led_matrix[i][j] = gaussian_spectrum(wl, &fitter->leds[i]);
        }
    }
}

float evaluate_fit(const SpectrumFitter* fitter, const float weights[], float* r2) {
    float combined[NUM_WAVEPOINTS] = {0};
    float target_mean = 0.0f;
    float sse = 0.0f, sst = 0.0f;
    
    // 组合光谱计算
    for (int i = 0; i < fitter->num_leds; i++) {
        if (weights[i] < WEIGHT_THRESHOLD) continue;
        for (int j = 0; j < NUM_WAVEPOINTS; j++) {
            combined[j] += weights[i] * fitter->led_matrix[i][j];
        }
    }
    
    // 计算统计量
    for (int j = 0; j < NUM_WAVEPOINTS; j++) {
        target_mean += fitter->target[j];
    }
    target_mean /= NUM_WAVEPOINTS;
    
    for (int j = 0; j < NUM_WAVEPOINTS; j++) {
        float err = combined[j] - fitter->target[j];
        sse += err * err;
        float dev = fitter->target[j] - target_mean;
        sst += dev * dev;
    }
    
    *r2 = (sst > EPSILON) ? (1.0f - sse/sst) : 0.0f;
    return sse / NUM_WAVEPOINTS;
}

void optimize_weights(SpectrumFitter* fitter) {
    float weights[MAX_LEDS];
    for (int i = 0; i < fitter->num_leds; i++) {
        weights[i] = 1.0f / fitter->num_leds;
    }
    
    for (int iter = 0; iter < MAX_ITER; iter++) {
        float r2, mse = evaluate_fit(fitter, weights, &r2);
        if (r2 >= R2_THRESHOLD) break;
        
        float grad[MAX_LEDS] = {0};
        const float eps = 0.001f;
        
        // 数值梯度计算
        for (int i = 0; i < fitter->num_leds; i++) {
            float temp = weights[i];
            weights[i] += eps;
            float new_mse = evaluate_fit(fitter, weights, &r2);
            float raw_grad = (new_mse - mse) / eps;
            grad[i] = fmaxf(fminf(raw_grad, GRAD_CLIP), -GRAD_CLIP);
            weights[i] = temp;
        }
        
        // 权重更新
        float sum = 0.0f;
        for (int i = 0; i < fitter->num_leds; i++) {
            weights[i] -= LEARNING_RATE * grad[i];
            weights[i] = fmaxf(weights[i], 0.0f);
            sum += weights[i];
        }
        
        // 归一化处理
        if (sum > EPSILON) {
            float inv_sum = 1.0f / sum;
            for (int i = 0; i < fitter->num_leds; i++) {
                weights[i] *= inv_sum;
            }
        } else {
            for (int i = 0; i < fitter->num_leds; i++) {
                weights[i] = 1.0f / fitter->num_leds;
            }
        }
    }
    
    memcpy(fitter->optimal_weights, weights, sizeof(weights));
}

void print_results(const SpectrumFitter* fitter) {
    float r2, mse;
    mse = evaluate_fit(fitter, fitter->optimal_weights, &r2);
    
    printf("=== Optimization Results ===\n");
    printf("R²: %.4f (Target ≥%.2f)\n", r2, R2_THRESHOLD);
    printf("MSE: %.6f\n", mse);
    
    float total_cost = 0.0f;
    printf("\nUsed LEDs:\n");
    for (int i = 0; i < fitter->num_leds; i++) {
        if (fitter->optimal_weights[i] > WEIGHT_THRESHOLD) {
            float cost = fitter->optimal_weights[i] * fitter->leds[i].cost;
            total_cost += cost;
            printf("  %2d: %4.0fnm (FWHM=%3.0f) | Weight=%.2f | Cost=%.1f\n",
                  i, fitter->leds[i].peak_wavelength, fitter->leds[i].fwhm,
                  fitter->optimal_weights[i], cost);
        }
    }
    printf("\nTotal Cost: %.1f\n", total_cost);
}

int main() {
    SpectrumFitter fitter;
    init_spectrum_fitter(&fitter);
    optimize_weights(&fitter);
    print_results(&fitter);
    return 0;
}